import ROOT
import numpy as np
from sklearn.cluster import Birch
import matplotlib.pyplot as plt
import sys
from scipy.spatial.distance import euclidean
from scipy.spatial import distance_matrix


np.set_printoptions(suppress=True, threshold=np.inf)
v = np.array([0, 0, -1])
I = np.identity(3)
INTER_PRECISION = 10
pic_count = 0

def calculate_vector_product(vector1, vector2):
    cross_product = (vector1[1] * vector2[2] - vector1[2] * vector2[1],
                     vector1[2] * vector2[0] - vector1[0] * vector2[2],
                     vector1[0] * vector2[1] - vector1[1] * vector2[0])
    return cross_product


def calculate_cos_sin_with_z_axis(f_vector):
    f_vector = f_vector / np.linalg.norm(f_vector)
    f_cos_theta = np.dot(f_vector, v)
    # print("f_cos_theta == " + str(f_cos_theta))
    f_sin_theta =  np.sqrt(1 - f_cos_theta ** 2)
    # print("f_sin_theta ==" + str(f_sin_theta))
    return f_cos_theta, f_sin_theta


def rotation_vector_parallel_z(f_array, f_start_point, f_end_point):
    f_vector = f_end_point - f_start_point
    # print("f_vector :", f_vector)
  
    f_vector = f_vector / np.linalg.norm(f_vector)
    # print("f_vector post :", f_vector)
  
    n = calculate_vector_product(f_vector, v)
    # print("n : ", n)
    n = n / np.linalg.norm(n)
    # print("n : ", n)
    N = np.array([
        [0, -n[2], n[1]],
        [n[2], 0, -n[0]],
        [-n[1], n[0], 0]
    ])
    # print("N:", N)
    N2 = np.dot(N, N)
    # print("N2:", N2)
    # cos sin
    f_cos_theta, f_sin_theta = calculate_cos_sin_with_z_axis(f_vector)

    R = I + f_sin_theta * N + (1 - f_cos_theta) * N2
    # print("R : ", R)
    f_array = np.dot(R, f_array.T)

    return f_array.T


def calculate_charge_center(charge_data):
    total_charge = np.sum(charge_data[:, 3])
    
    charge_center_x = np.sum(charge_data[:, 0] * charge_data[:, 3].T) / total_charge
    charge_center_y = np.sum(charge_data[:, 1] * charge_data[:, 3].T) / total_charge
    charge_center_z = np.sum(charge_data[:, 2] * charge_data[:, 3].T) / total_charge
    charge_center = np.array([charge_center_x, charge_center_y, charge_center_z])
    # print('charge_center:', charge_center)
    return charge_center



# Root data
root_file = ROOT.TFile('/lustre/pandax/wangxuan/restTest/MeGat/v2_cgc/GasCZTCombine_e_z-_wxenergyMeV_n1w.root', 'READ')
MLTree = root_file.CombineHits
output_file = ROOT.TFile('/lustre/pandax/wangxuan/restTest/DataEnhance/v2_pydata/KFData_e_z-_wxenergyMeV_n1w.root', 'RECREATE')
KFData = ROOT.TTree('KFData', 'Data for KF')

eventID = ROOT.TArrayD()
enhance_data_fx = ROOT.TArrayD()
enhance_data_fy = ROOT.TArrayD()
enhance_data_fz = ROOT.TArrayD()
G4_e_vertex_x = ROOT.TArrayD()
G4_e_vertex_y = ROOT.TArrayD()
G4_e_vertex_z = ROOT.TArrayD()
gamma_x = ROOT.TArrayD()
gamma_y = ROOT.TArrayD()
gamma_z = ROOT.TArrayD()
electron_energy = ROOT.TArrayD()
gasDepositedEnergy = ROOT.TArrayD()
totalDepositedEnergy_smear = ROOT.TArrayD()

KFData.Branch("id", eventID)
KFData.Branch("fx", enhance_data_fx)
KFData.Branch("fy", enhance_data_fy)
KFData.Branch("fz", enhance_data_fz)
KFData.Branch("G4_e_vertex_x", G4_e_vertex_x)
KFData.Branch("G4_e_vertex_y", G4_e_vertex_y)
KFData.Branch("G4_e_vertex_z", G4_e_vertex_z)
KFData.Branch("gamma_x", gamma_x)
KFData.Branch("gamma_y", gamma_y)
KFData.Branch("gamma_z", gamma_z)
KFData.Branch("electron_energy", electron_energy)
KFData.Branch("gasDepositedEnergy", gasDepositedEnergy)
KFData.Branch("totalDepositedEnergy_smear", totalDepositedEnergy_smear)
count = 0
for entry in MLTree:
    # count += 1
    if count<20:
    # if count<20 and entry.id == 10:
        eventID.Set(1, np.array([entry.id], dtype=np.float64))
        dx = entry.dx
        dy = entry.dy
        dz = entry.dz
        if dx<50 and dy<50 and dz<50:
            continue
        G4_e_vertex_x.Set(1, np.array([entry.G4_e_vertex_x], dtype=np.float64))
        G4_e_vertex_y.Set(1, np.array([entry.G4_e_vertex_y], dtype=np.float64))
        G4_e_vertex_z.Set(1, np.array([entry.G4_e_vertex_z], dtype=np.float64))
        gamma_x.Set(1, np.array([entry.gamma_x], dtype=np.float64))
        gamma_y.Set(1, np.array([entry.gamma_y], dtype=np.float64))
        gamma_z.Set(1, np.array([entry.gamma_z], dtype=np.float64))
        fx = entry.gas_x
        fy = entry.gas_y
        fz = entry.gas_z
        fe = entry.gas_e
        electron_energy.Set(1, np.array([entry.electron_energy], dtype=np.float64))
        gasDepositedEnergy.Set(1, np.array([entry.gasDepositedEnergy], dtype=np.float64))
        totalDepositedEnergy_smear.Set(3, np.array([entry.totalDepositedEnergy_smear], dtype=np.float64))

        data = np.column_stack((fx, fy, fz, fe))

        # Create and fit the Birch model
        birch = Birch(threshold=20, n_clusters=None)
        birch.fit(data[:, :3])
        cluster_labels = birch.labels_
        data = np.insert(data, data.shape[1], cluster_labels.T, axis=1)
        # print('cluster hits data:', data)
        
        # 按label分组创建字典
        cluser_dict = {}
        for arr in data:
            key = arr[4]
            if key in cluser_dict:
                cluser_dict[key].append(arr[:4])
            else:
                cluser_dict[key] = [arr[:4]] 

        cluster_charge_center = []
        for key, value in cluser_dict.items():
            value = np.array(value)
            charge_center = calculate_charge_center(value)
            cluster_charge_center.append(charge_center)
            # print(f"The charge center for key {key} is {charge_center}")
        
        cluster_centers = np.array(cluster_charge_center)
        # cluster_centers = birch.subcluster_centers_
        # cluster_centers = cluster_centers[1:]
        # cluster_centers = np.insert(cluster_centers, 0, [fx[0], fy[0], fz[0]], axis=0)
        
        # Get the cluster labels
        

        # Print the cluster labels
        cluster_centers = np.round(cluster_centers, 4)
        # print(cluster_centers)
        # print(cluster_centers.shape[0])
        if (cluster_centers.shape[0]) < 4:
            continue

        # Sort by distance
        dist_matrix = distance_matrix(cluster_centers, cluster_centers)
        visited = [0]
        current_point = 0
        sorted_points = [cluster_centers[current_point]]
        
        while len(visited) < len(cluster_centers):
            min_distance = np.inf
            next_point = None
            
            for i in range(len(cluster_centers)):
                if i not in visited and dist_matrix[current_point, i] < min_distance:
                    min_distance = dist_matrix[current_point, i]
                    next_point = i
            
            visited.append(next_point)
            current_point = next_point
            sorted_points.append(cluster_centers[next_point])
        
        new_points = sorted_points[0] - 2 * (sorted_points[1] - sorted_points[0])
        sorted_points = np.insert(sorted_points, 0, new_points, axis=0)
        # for idx, point in enumerate(sorted_points):
        #     print(f"Point {idx + 1}: {point}")

        interpolated_points = []
        interpolated_points.append(sorted_points[0])
        interpolated_points = np.array(interpolated_points)
        interpolated_points = np.insert(interpolated_points, interpolated_points.shape[1], 0, axis=1)

        # Interpolation
        current_point = sorted_points[0]
        correct_distance = 0
        type_count = 0
        for next_point in sorted_points[1:]:
            new_cluster_interpolated_points = []
            step = next_point - current_point
            step = step / np.linalg.norm(step)
            distance = euclidean(current_point, next_point)
            if distance>INTER_PRECISION:
                interpolated_point = current_point + step * (INTER_PRECISION - correct_distance)
                correct_distance = 0
                new_cluster_interpolated_points.append(interpolated_point)

                distance = euclidean(new_cluster_interpolated_points[-1], next_point)
                num_interpolated_points = int(distance / INTER_PRECISION)

                for i in range(1, num_interpolated_points + 1):
                    interpolated_point = new_cluster_interpolated_points[-1] + step * INTER_PRECISION
                    new_cluster_interpolated_points.append(interpolated_point)
        
                correct_distance += euclidean(new_cluster_interpolated_points[-1], next_point)
                new_cluster_interpolated_points.append(next_point)
                new_cluster_interpolated_points = np.array(new_cluster_interpolated_points)

                cluster_type = np.ones((new_cluster_interpolated_points.shape[0], 1), dtype=int) * type_count
                new_cluster_interpolated_points = np.insert(new_cluster_interpolated_points, new_cluster_interpolated_points.shape[1], cluster_type.T, axis=1)
                type_count += 1
                current_point = next_point
                interpolated_points = np.row_stack((interpolated_points, new_cluster_interpolated_points))

        
        # for idx, point in enumerate(interpolated_points[:, :3]):
        #     if idx > 0:
        #         distance = euclidean(point, interpolated_points[idx - 1, :3])
        #         print(f"Interpolated Point {idx}: {point}, Distance to Previous: {distance:.2f}")
        #     else:
        #         print(f"Interpolated Point {idx}: {point}, Distance to Previous: 0.00")
        # print('interpolated_points with labels', interpolated_points)


        #_________calculate new interpolation points with charge center_______________
        charge_center_interpolation_points = []
        for idx in range(interpolated_points.shape[0]-1):
            cluster_hits = []
            selected_points = []
            cluster_type_id = interpolated_points[idx, 3]
            # print('cluster_type_id', cluster_type_id)
    
            # If the two selected slice points are within the same cluster,
            # select this cluster. Otherwise, still need to join the next cluster
            if interpolated_points[idx, 3] == interpolated_points[idx+1, 3]: 
                for hits in data:
                    if hits[4] == cluster_type_id:
                        cluster_hits.append(hits)
            else:
                for hits in data:
                    if hits[4] == cluster_type_id or hits[4] == cluster_type_id+1:
                        cluster_hits.append(hits)

            cluster_hits = np.array(cluster_hits)
            if np.any(cluster_hits):
                # print('data for calculate charge center_', idx, ': ', cluster_hits)
                pre_start = interpolated_points[idx, :3]
                pre_end = interpolated_points[idx+1, :3]
                # print('pre_start_', idx, ': ', pre_start)
                # print('pre_end', idx, ': ', pre_end)
                transformed_electron_data = rotation_vector_parallel_z(cluster_hits[:, :3], pre_start, pre_end)
                new_start = rotation_vector_parallel_z(interpolated_points[idx, :3], pre_start, pre_end)
                new_end = rotation_vector_parallel_z(interpolated_points[idx+1, :3], pre_start, pre_end)
            
                for i in range(transformed_electron_data.shape[0]):
                    if new_end[2] < transformed_electron_data[i, 2] <= new_start[2]:
                        selected_points.append(cluster_hits[i, :4])
                selected_points = np.array(selected_points)
                if np.any(selected_points):
                    # print('selected_points: ', selected_points)
                    total_charge = np.sum(selected_points[:, 3])
                    if selected_points.shape[0] > 5 and total_charge!= 0:
                        charge_center = calculate_charge_center(selected_points)
                        # print('charge_center: ', charge_center)
                        charge_center_interpolation_points.append(charge_center)

                # else:
                    # print('no selected_points')
                    # print('cluster hits: ', cluster_hits[:, :3])
                    # print('pre start and end:', pre_start, pre_end)
                    # print('transformed hits: ', transformed_electron_data)
                    # print('new start and end:', new_start, new_end)
                
        charge_center_interpolation_points = np.array(charge_center_interpolation_points)
        # sorted_points = np.array(sorted_points)
        # if np.any(charge_center_interpolation_points):
        #     enhance_data_points = np.row_stack((charge_center_interpolation_points, sorted_points[:, :3]))
        # else:
        #     continue
        # # Sort by distance
        # enhance_data_points = np.array(enhance_data_points)
        # dist_matrix_data = distance_matrix(enhance_data_points, enhance_data_points)
        # visited = [0]
        # current_point_data = 0
        # sorted_data_points = [enhance_data_points[current_point_data]]

        # while len(visited) < len(enhance_data_points):
        #     min_distance = np.inf
        #     next_point = None
            
        #     for k in range(len(enhance_data_points)):
        #         if k not in visited and dist_matrix_data[current_point_data, k] < min_distance:
        #             min_distance = dist_matrix_data[current_point_data, k]
        #             next_point = k
            
        #     visited.append(next_point)
        #     current_point_data = next_point
        #     sorted_data_points.append(enhance_data_points[next_point])
        # sorted_data_points = np.array(sorted_data_points)

        # combined_coords = np.array(interpolated_points, dtype=np.float64)
        if charge_center_interpolation_points.shape[0]<5:
            continue
        
        tmp_fx = np.array(charge_center_interpolation_points[:, 0], dtype=np.float64)
        tmp_fy = np.array(charge_center_interpolation_points[:, 1], dtype=np.float64)
        tmp_fz = np.array(charge_center_interpolation_points[:, 2], dtype=np.float64)
        array_size = len(tmp_fx)
        enhance_data_fx.Set(array_size, tmp_fx)
        enhance_data_fy.Set(array_size, tmp_fy)
        enhance_data_fz.Set(array_size, tmp_fz)

        KFData.Fill()
        if entry.id % 100 == 0 :
            print('process entry_%s' % entry.id)
        # _____Plot the cluster labels
        
        if pic_count < 20:
            fig = plt.figure(figsize=(8, 8))

            # select cluser 0 and 1
            # for k in range(data.shape[0]):
            #     if data[k][4] == 0:    
            #         plt.scatter(data[k][0], data[k][2], c=0, s=20)
            #     # if data[k][4] == 1:
            #     #     plt.scatter(data[k][0], data[k][2], c=10, s=20)
            # plt.scatter(cluster_centers[0][0], cluster_centers[0][2], c='red', s=30)
            # plt.scatter(cluster_centers[1][0], cluster_centers[1][2], c='blue', s=30)

            plt.scatter(G4_e_vertex_x, G4_e_vertex_z, c='blue',s=30)
            plt.scatter(data[:, 0], data[:, 2], c=cluster_labels, s=20)
            plt.scatter(tmp_fx, tmp_fz, c='red', s=30)
            plt.scatter(interpolated_points[:, 0], interpolated_points[:, 2], marker='*', c='green', s=30)

            plt.scatter(cluster_centers[:, 0], cluster_centers[:, 2], marker='*', c='blue', s=50)
            # plt.xlim(-150, 150)
            # plt.ylim(-150, 150) 
            plt.xlabel('x[mm]')
            plt.ylabel('z[mm]')
            plt.title('Cluster Centers')
            plt.savefig('/lustre/pandax/wangxuan/restTest/DataEnhance/pic_wxenergyMeV_5mm/CC_intpol_%s.png' % entry.id)
            pic_count += 1

output_file.Write()
output_file.Close()
print("KF data has been stored in root.")
